import matplotlib.pyplot as plt
import torch
from PIL import Image

from config import Config
from model import AlexNet
from tools import getTransform


def cvtImg(im):
    im_r = im.resize((32, 32))
    im_v = transform(im_r)
    im_v = torch.unsqueeze(im_v, dim=0)
    im_v = im_v.to(device)
    return im_v


# get transform
transform = getTransform()
# get device which will be used to load network
device = torch.device('cuda:0' if torch.cuda.is_available else 'cpu')
# create a network
pred_net = AlexNet(len(Config.CONFIG_CLASSES))
# load weights
pred_net.load_state_dict(torch.load(Config.CONFIG_SAVE_PATH))
# move to device
pred_net.to(device)
# open image
im = Image.open(Config.CONFIG_TEST_IMG)
# convert some properties
im_v = cvtImg(im)

pred_net.eval()
# close auto grad function using context no_grad
with torch.no_grad():
    output = pred_net(im_v)
    predict = torch.softmax(output, dim=1)
    label = Config.CONFIG_CLASSES[predict[0].argmax().item()]
    plt.imshow(im)
    plt.text(im.size[0] / 2, -im.size[1] / 50, label, ha='center', fontsize=12)
    plt.show()
